// loop_control_cpp/src/path_controller.cpp
#include <rclcpp/rclcpp.hpp>
#include <geometry_msgs/msg/twist.hpp>
#include <nav_msgs/msg/odometry.hpp>
#include <tf2/LinearMath/Quaternion.h>
#include <tf2/LinearMath/Matrix3x3.h>
#include <memory>
#include <vector>
#include <cmath>

using namespace std::chrono_literals;

class PathController : public rclcpp::Node {
public:
    PathController() : Node("path_controller") {
        // Initialize parameters
        declare_parameter("kp_linear", 0.5);
        declare_parameter("ki_linear", 0.01);
        declare_parameter("kd_linear", 0.1);
        declare_parameter("kp_angular", 0.8);
        declare_parameter("ki_angular", 0.01);
        declare_parameter("kd_angular", 0.1);
        declare_parameter("max_linear_speed", 0.5);
        declare_parameter("max_angular_speed", 1.0);

        // Target points sequence
        target_points_ = {
            {1.0, 0.0},
            {5.0, 1.0},
            {5.0, 4.0},
            {1.0, 4.0}
        };

        // ROS interfaces
        cmd_vel_pub_ = create_publisher<geometry_msgs::msg::Twist>("/cmd_vel", 10);
        odom_sub_ = create_subscription<nav_msgs::msg::Odometry>(
            "/odom", 10, std::bind(&PathController::odom_callback, this, std::placeholders::_1));

        control_timer_ = create_wall_timer(
            100ms, std::bind(&PathController::control_loop, this));

        // Initial state
        state_ = State::ROTATING;
        RCLCPP_INFO(get_logger(), "Path controller initialized!");
    }

private:
    struct Point {
        double x;
        double y;
    };

    enum class State { MOVING, ROTATING, FINISHED };

    void odom_callback(const nav_msgs::msg::Odometry::SharedPtr msg) {
        current_pose_ = msg->pose.pose;
    }

    double quaternion_to_yaw(const geometry_msgs::msg::Quaternion& q) {
        tf2::Quaternion tf_q(q.x, q.y, q.z, q.w);
        tf2::Matrix3x3 m(tf_q);
        double roll, pitch, yaw;
        m.getRPY(roll, pitch, yaw);
        return yaw;
    }

    void control_loop() {
        if (!current_pose_ || state_ == State::FINISHED) {
            return;
        }

        auto [distance_error, angle_error] = calculate_errors();

        // State transition logic
        handle_state_transition(distance_error, angle_error);

        // Generate control commands
        auto cmd_vel = geometry_msgs::msg::Twist();
        if (state_ == State::MOVING) {
            // Linear control
            double linear_speed = pid_control(
                distance_error, 
                last_linear_error_, 
                integral_linear_,
                get_parameter("kp_linear").as_double(),
                get_parameter("ki_linear").as_double(),
                get_parameter("kd_linear").as_double()
            );
            cmd_vel.linear.x = std::clamp(linear_speed, 
                -get_parameter("max_linear_speed").as_double(),
                get_parameter("max_linear_speed").as_double());

            // Angular control
            double angular_speed = pid_control(
                angle_error, 
                last_angular_error_, 
                integral_angular_,
                get_parameter("kp_angular").as_double(),
                get_parameter("ki_angular").as_double(),
                get_parameter("kd_angular").as_double()
            );
            cmd_vel.angular.z = std::clamp(angular_speed,
                -get_parameter("max_angular_speed").as_double(),
                get_parameter("max_angular_speed").as_double());
        }
        else if (state_ == State::ROTATING) {
            // Only angular control when rotating
            double angular_speed = pid_control(
                angle_error, 
                last_angular_error_, 
                integral_angular_,
                get_parameter("kp_angular").as_double(),
                get_parameter("ki_angular").as_double(),
                get_parameter("kd_angular").as_double()
            );
            cmd_vel.angular.z = std::clamp(angular_speed,
                -get_parameter("max_angular_speed").as_double(),
                get_parameter("max_angular_speed").as_double());
        }

        cmd_vel_pub_->publish(cmd_vel);
    }

    std::pair<double, double> calculate_errors() {
        double current_x = current_pose_->position.x;
        double current_y = current_pose_->position.y;
        const auto& target = target_points_[current_target_];

        double dx = target.x - current_x;
        double dy = target.y - current_y;
        double distance_error = std::hypot(dx, dy);

        double target_yaw = std::atan2(dy, dx);
        double current_yaw = quaternion_to_yaw(current_pose_->orientation);
        double angle_error = target_yaw - current_yaw;

        // Normalize angle error to [-π, π]
        angle_error = std::atan2(std::sin(angle_error), std::cos(angle_error));

        return {distance_error, angle_error};
    }

    double pid_control(double error, double& last_error, double& integral,
                      double kp, double ki, double kd) {
        // Integral term with anti-windup
        integral = std::clamp(integral + error * 0.1, 
            -10.0, 10.0); // 100ms period
        
        // Derivative term
        double derivative = (error - last_error) / 0.1;
        last_error = error;

        return kp * error + ki * integral + kd * derivative;
    }

    void handle_state_transition(double distance_error, double angle_error) {
        if (state_ == State::FINISHED) return;

        if (state_ == State::ROTATING) {
            // Tighter angle threshold for better accuracy
            if (std::abs(angle_error) < 0.05) {
                state_ = State::MOVING;
                // Reset linear PID terms when transitioning to moving
                integral_linear_ = 0.0;
                last_linear_error_ = 0.0;
                RCLCPP_INFO(get_logger(), "Orientation achieved, moving to target %d", current_target_);
            }
        }
        else if (state_ == State::MOVING) {
            if (distance_error < 0.1) {
                if (++current_target_ >= target_points_.size()) {
                    state_ = State::FINISHED;
                    stop_robot();
                    RCLCPP_INFO(get_logger(), "All targets reached!");
                }
                else {
                    state_ = State::ROTATING;
                    // Reset angular PID terms when transitioning to rotating
                    integral_angular_ = 0.0;
                    last_angular_error_ = 0.0;
                    RCLCPP_INFO(get_logger(), "Reached target %d", current_target_-1);
                }
            }
        }
    }

    void stop_robot() {
        auto cmd_vel = geometry_msgs::msg::Twist();
        cmd_vel_pub_->publish(cmd_vel);
    }

    // Configuration
    std::vector<Point> target_points_;
    size_t current_target_ = 0;
    State state_;

    // Current robot state
    std::optional<geometry_msgs::msg::Pose> current_pose_;

    // PID control state
    double last_linear_error_ = 0.0;
    double integral_linear_ = 0.0;
    double last_angular_error_ = 0.0;
    double integral_angular_ = 0.0;

    // ROS interfaces
    rclcpp::Publisher<geometry_msgs::msg::Twist>::SharedPtr cmd_vel_pub_;
    rclcpp::Subscription<nav_msgs::msg::Odometry>::SharedPtr odom_sub_;
    rclcpp::TimerBase::SharedPtr control_timer_;
};

int main(int argc, char** argv) {
    rclcpp::init(argc, argv);
    auto node = std::make_shared<PathController>();
    rclcpp::spin(node);
    rclcpp::shutdown();
    return 0;
}